import argparse
import json
import os

from src.log.logger import logger
from src.lib.sbom_generate import SbomGenerate
from src.tools.query import RpmInfoQuery
from config.constants import SBOM_OUTPUT_PATH


def main(args):
    rpm_path =  args.rpm_path
    spec_path = args.spec_path
    output_path = args.output_path
    if args.output_path is None:
        output_path = SBOM_OUTPUT_PATH
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    logger.info("=======Input rpm_path: %s, spec_path: %s, output_path: %s", rpm_path, spec_path, output_path)
    
    logger.info("======Start get rpms provides and depends....")
    rpm_files = RpmInfoQuery.find_files(rpm_path, 'rpm')
    rpms_provides = RpmInfoQuery.query_rpms_provides(rpm_files)
    rpms_depends = RpmInfoQuery.query_rpms_depends(rpm_files)
    rpms_infos = RpmInfoQuery.query_rpms_info(rpm_files)
    rpm_repo_data = { "rpms_provides": rpms_provides, "rpms_depends" : rpms_depends, 'rpms_infos': rpms_infos }
    with open(os.path.join(output_path, 'rpm_repo_data.json'), 'w') as output:
        json.dump(rpm_repo_data, output, indent=2)
    logger.info("======End get rpms provides and depends")
    
    for rpm_file in rpm_files:
        logger.info("======start processing rpm_file: %s", rpm_file)
        sbom_generator=SbomGenerate(rpm_file, spec_path, output_path)
        top_element_data = sbom_generator.generate_sbom_top()
        logger.info("======Got top_data")
        packages_element_data = sbom_generator.generate_sbom_packages(rpms_infos.get(rpm_file, {}))
        logger.info("======Got packages_data")
        files_element_data = sbom_generator.generate_sbom_files()
        logger.info("======Got files_element_data")
        relationships_element_data = sbom_generator.generate_sbom_relationships(rpms_depends.get(rpm_file, []), \
                                                                                rpms_infos, \
                                                                                rpms_provides)
        logger.info("======Got relationships_element_data")
        merge_files = {
            "top":top_element_data,
            "packages": packages_element_data,
            "files": files_element_data,
            "relationships": relationships_element_data
        }
        
        sbom_generator.merge_sbom_element_jsons(merge_files)
        

if __name__ == "__main__":
    # 创建一个 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='Process sbom info of rpms.')

    # 添加命令行参数
    parser.add_argument('-r', '--rpm_path', type=str, required=True, help='The rpm path to process.')
    parser.add_argument('-s', '--spec_path', type=str, required=True, help='The spec path to process.')
    parser.add_argument('-o', '--output_path', type=str, default=None, help='The output path to process.')

    # 解析命令行参数
    args = parser.parse_args()

    # 使用解析后的参数
    main(args)